import pandas
import numpy
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.cluster import KMeans
from mpl_toolkits.mplot3d import Axes3D
from mpl_toolkits.mplot3d import proj3d

wine = pandas.read_csv('wine.csv')
pandas.set_option('display.width', None)
# print(wine.head())

X = wine.drop('Class', axis=1)
Y = wine.Class

x_tr, x_t, y_tr, y_t = train_test_split(X, Y, train_size=.80)

transfer = StandardScaler()
x_tr = transfer.fit_transform(x_tr)
x_t = transfer.fit_transform(x_t)

model = KNeighborsClassifier()
model.fit(x_tr, y_tr)

pre = model.predict(x_t)
print("Prediction accuracy: " + str(model.score(x_t, y_t)))

estimator = KMeans(n_clusters=3)

Y = estimator.fit_predict(X)

T = numpy.array(X)
fig = plt.figure(figsize=(5, 5))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(T[:, 0], T[:, 1], T[:, 2], c=Y, marker='*')
plt.show()
